package study.concurrent.thread.support.utils;

import org.slf4j.MDC;

import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.Executor;
import java.util.function.Supplier;

// https://stackoverflow.com/questions/55899284/mdc-logger-with-completablefuture
public class AsyncUtils {

    public static <U> CompletableFuture<U> supplyAsync(Supplier<U> supplier) {
        return CompletableFuture.supplyAsync(new SupplierWrapper(supplier));
    }

    public static <U> CompletableFuture<U> supplyAsync(Supplier<U> supplier, Executor executor) {
        return CompletableFuture.supplyAsync(new SupplierWrapper(supplier), executor);
    }

    //参考：https://qa.1r1g.com/sf/ask/3437752181/
    private static class SupplierWrapper<T> implements Supplier<T> {

        private final Supplier<T> delegate;

        private final Map<String, String> mdc;

        public SupplierWrapper(Supplier<T> delegate) {
            this.delegate = delegate;
            this.mdc = MDC.getCopyOfContextMap();
        }

        @Override
        public T get() {
            MDC.setContextMap(mdc);
            return delegate.get();
        }
    }
}
